// decode.c: Lightweight HTTP/2 HPACK Huffman decompressor.
// Written 5/30/14 by David Krauss. No rights reserved.

#include <stdint.h>
#include <limits.h>
#include <stddef.h>
#include <stdbool.h>

#if TRACE_HPACK_HUFFMAN
#include <stdio.h>
#endif

#include "decode.h"


// These primitives must be defined and linked to provide this library with a data source.
// The bundled bitstream.c/h files are intended as a reference, hence the header is not #included.
struct hpack_huffman_bit_source;

// Equivalent to get_bits( source, 0, 8 ).
extern uint8_t hpack_huffman_get_byte( struct hpack_huffman_bit_source * source );

// Shift the next n bits from source into head. If no more bits exist, return UCHAR_MAX.
extern uint8_t hpack_huffman_get_bits( struct hpack_huffman_bit_source * source, uint8_t head, uint8_t n );

// Restore the last n bits from tail into the source. This library will not unget bits it did not get.
extern void hpack_huffman_unget_bits( struct hpack_huffman_bit_source * source, uint8_t tail, uint8_t n );

// Return whether the source is done.
extern bool hpack_huffman_at_end( struct hpack_huffman_bit_source * source );


/*	These tables of values record the maximum value of the code for each length, with the 
	code MSB aligned to the decode register MSB.
	
	When the code length exceeds one byte, the remaining codes all have 1's in the first
	N most significant bits. The algorithm right-shifts those bits away, and shifts in new
	bits. The next section of the table records the maximum value of each shifted coding
	at each given length. */

enum code_breaks_e {
	// First 8 bits. All the actual compression happens here.
	cb4 = 0x30,
	cb5 = 0x78,
	cb6 = 0xC8,
	cb7 = 0xEC, 
	cb8 = 0xF8,
	
	// Shift 5 bits, interpret up to 13.
	cb9 = 0xD0,
	cb10 = 0xF0,
	cb11 = 0xF8,
	cb12 = 0xFC,
	cb13 = 0xFE,
	
	// Shift 7 bits, interpret up to 20.
	cb14 = 0x80,
	cb15 = 0xE0,
	cb16 = 0xE0,
	cb17 = 0xF8,
	cb18 = 0xFC, // After 18, remaining codes are 25-26 bits.
	cb19 = 0xFC,
	cb20 = 0xFC,
	
	// Shift 6 bits, interpret up to 26.
	cb21 = 0, cb22 = 0, cb23 = 0, cb24 = 0,
	cb25 = 0xBA,
	cb26 = 0x100
};

static uint8_t const segment_sizes[] = { /* 4, */ 4, 6, 5 };

/*	Encode the differences between the codes, so that they may be progressively subtracted
	from the register. When the register is less than the next difference, we know the
	length of the current code. */
	
static uint8_t const code_breaks[] = {
	cb4, // 30, 48, 50, 24, 0C
	cb5 - cb4,
	cb6 - cb5,
	cb7 - cb6,
	cb8 - cb7,
	
	cb9, // D0, 20, 8, 4, 2
	cb10 - cb9,
	cb11 - cb10,
	cb12 - cb11,
	cb13 - cb12,
	
	cb14, // 80, 70, 0, 18, 4, 0, 0
	cb15 - cb14,
	cb16 - cb15,
	cb17 - cb16,
	cb18 - cb17,
	cb19 - cb18,
	cb20 - cb19,
	
	cb21, // 0, 0, 0, 0, BA, 46
	cb22 - cb21,
	cb23 - cb22,
	cb24 - cb23,
	cb25 - cb24,
	cb26 - cb25
};

/*	All the symbols with encodings less than the default "huge" length.
	Sorted in order of encoding length, then lexicographic code order. */

static char const decode_table[] = // 94 characters: printing ASCII, except backslash is special case.
	"012 /3aceiot%.45"
	"6789:=Tdghlmnprs"
	"&,-ADFGMNS_bfuvw"
	"xy;BCEIOPUXjkz\")"
	"HJKLQRVWYZq(*+?["
	"]|~!'#^$>@<{}`";

static int hpack_huffman_get_char( struct hpack_huffman_bit_source * in_bits ) {
	uint8_t const * segment_size = segment_sizes;
	uint8_t const * break_diff = code_breaks;
	
	uint8_t out_index = 0;
	
	uint8_t tail_discard_count = 4;
	uint8_t in = hpack_huffman_get_byte( in_bits );
#if TRACE_HPACK_HUFFMAN
	printf( "decode byte %x\n", in );
#endif
	
	for (;;) {
		do {
			if ( in < * break_diff ) {
				out_index += in >> tail_discard_count;
				goto do_index;
			}
			
			out_index += * break_diff >> tail_discard_count;
			in -= * break_diff;
			
			++ break_diff;
			
		} while ( tail_discard_count -- );
		
		tail_discard_count = * segment_size ++;
		in = hpack_huffman_get_bits( in_bits, in, tail_discard_count + 1 );

#if TRACE_HPACK_HUFFMAN
		printf( "decode segment %x index = %d\n", in, out_index );
#endif
		
		if ( in == UCHAR_MAX && hpack_huffman_at_end( in_bits ) ) return -1;
	}
	
do_index:
	
	hpack_huffman_unget_bits( in_bits, in, tail_discard_count );
	
	if ( out_index < sizeof decode_table - 1 ) {
		return decode_table[ out_index ];
	} else {
		uint8_t const code_count_25bit = cb25 / 2;
		uint8_t const backslash_index = code_count_25bit + 32;
	
		out_index -= sizeof decode_table - 1;
		
#if TRACE_HPACK_HUFFMAN
		printf( "special index %d vs BS = %d\n", out_index, backslash_index );
#endif
		
		if ( out_index < code_count_25bit ) {
			return out_index + 164; // Return 25-bit code; this contiguous range starts at 164 = 0xA4.
		} else if ( out_index < backslash_index ) {
			return out_index - code_count_25bit + 0; // 26-bit codes from range starting at NUL = 0.
		} else if ( out_index == backslash_index ) {
			return '\\';
		} else { // 26-bit codes from range starting at DEL = 127.
			return out_index - code_count_25bit + 127 - 33 /* 32 codes starting from NUL, one for backslash */;
		}
	}
}

extern int hpack_huffman_decode( struct hpack_huffman_bit_source * in_bits, uint8_t ** out, size_t max ) {
	uint8_t * pen = * out;
	for ( uint8_t * out_max = pen + max; pen != out_max; ++ pen ) {
		int c = hpack_huffman_get_char( in_bits );
		
		if ( c == 256 ) { // A Huffman encoded string literal containing the EOS entry MUST be treated as a decoding error. 
			return hpack_huffman_decode_status_eos;
		} else if ( c == -1 ) {
			* out = pen;
			return hpack_huffman_decode_status_ok;
		} else {
#if TRACE_HPACK_HUFFMAN
			printf( "emit %c = %d\n", c, c );
#endif
			* pen = c;
		}
	}
	
	* out = pen;
	return hpack_huffman_decode_status_overrun;
}

